# pullVideo.py
import numpy

import asyncio
import websockets
import numpy as np
import cv2
import cv2 as cv
import base64
import time
# import mediapipe as mp
import sys
sys.path.insert(0, '/disk1/chengw/django_vue_demo/django-demo/yolov5_face/')
from yolov5 import YoloV5

# cv2.VideoCapture(0)为获取摄像头，因为电脑没有摄像头所以使用视频代替
# capture = cv2.VideoCapture(r'/disk1/test_video/264a.mp4')
capture = cv2.VideoCapture(r'rtsp://admin:admin@10.82.14.206/1/h264major')
detector = YoloV5()
# mpHands = mp.solutions.hands
# hands = mpHands.Hands()
# mpDraw = mp.solutions.drawing_utils
# # 设置点的颜色与大小
# handLmsStyle = mpDraw.DrawingSpec(color=(0, 0, 255), thickness=5)
# # 设置线的颜色与大小
# handConStyle = mpDraw.DrawingSpec(color=(0, 255, 0), thickness=10)
if not capture.isOpened():
    print('quit')
    quit()
encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 95]


# 开始推流视频
async def pullVideo(websocket):
    while True:
        ret, frame = capture.read()
        if ret:
            res = detector.inference_image(frame)
            # frame_RGB = cv.cvtColor(frame, cv.COLOR_BGR2RGB)  # 将RGB转为BGR
            # result = hands.process(frame_RGB)
            # if result.multi_hand_landmarks:
            #     for handLms in result.multi_hand_landmarks:
            #         mpDraw.draw_landmarks(
            #             frame, handLms, mpHands.HAND_CONNECTIONS, handLmsStyle, handConStyle)
        else:
            break
        # time.sleep(0.01)
        result, imgencode = cv2.imencode('.jpg', res, encode_param)
        data = np.array(imgencode)
        img = data.tobytes()
        # base64编码传输
        img = base64.b64encode(img).decode()
        await websocket.send("data:image/jpg;base64," + img)


async def main_logic():
    async with websockets.connect('ws://10.82.14.162:8000/ws/video/wms/', ping_interval=None) as websocket:
        await pullVideo(websocket)


asyncio.get_event_loop().run_until_complete(main_logic())
